import sys
from PyQt5.QtWidgets import (QWidget, QPushButton, QApplication,
                             QHBoxLayout, QVBoxLayout, QMessageBox,
                             QLabel)
from PyQt5 import QtGui, QtCore, QtWidgets
import pyqtgraph as pg
import tushare as ts
import datetime


# 通过股票代码获取公司名称
def getName(stoke_code):
    dat = ts.pro_api().query('stock_basic', fields='symbol,name')
    company_name = list(dat.loc[dat['symbol'] == stoke_code].name)[0]
    return company_name


# K线图表类
class CandlesBarChart(QWidget):
    def __init__(self, parent):
        super().__init__()
        self.parent = parent
        self.setWindowTitle("K线图")

        self.stock_code = ""  # 目前正在查询的股票代码
        self.query_type = "daily"  # 目前的查询波段(日,周,月)

        # 创建界面上的样式
        self.btnReture = QPushButton("返回列表")
        self.btnReture.clicked.connect(self.on_btnReture)
        self.stockLabel = QLabel("[000001 - 平安银行]")  # 显示股票信息的标签
        self.stockLabel.setStyleSheet("color:red")
        self.dayLabel = QLabel("天数:")
        self.daySel = QtWidgets.QComboBox()  # 天数下拉框
        self.daySel.addItem("近7天")
        self.daySel.addItem("近30天")
        self.daySel.addItem("近60天")
        self.daySel.addItem("近120天")
        self.daySel.addItem("近180天")
        self.daySel.addItem("近360天")
        self.daySel.addItem("近720天")
        self.daySel.setCurrentIndex(3)  # 默认显示120天的数据
        self.mainIndexLabel = QLabel("主图指标:")
        self.mainIndexSel = QtWidgets.QComboBox()  # 主图指标下拉框
        self.mainIndexSel.addItem("无")
        self.index1Label = QLabel("指标1:")
        self.index1Sel = QtWidgets.QComboBox()  # 主图指标下拉框
        self.index1Sel.addItem("无")
        self.index2Label = QLabel("指标2:")
        self.index2Sel = QtWidgets.QComboBox()  # 主图指标下拉框
        self.index2Sel.addItem("无")
        self.analyseLabel = QLabel("分析:")
        self.analyseSel = QtWidgets.QComboBox()
        self.analyseSel.addItem("不分析")
        self.rangeLabel = QLabel("波段:")
        self.rangeSel = QtWidgets.QComboBox()
        self.rangeSel.addItem("日线")
        self.rangeSel.addItem("周线")
        self.rangeSel.addItem("月线")

        # self.analyseSel.addItem("缠论")
        # self.analyseSel.addItem("波浪")
        self.btnQuery = QPushButton("查询")
        self.btnQuery.clicked.connect(self.on_btnQuery)
        hbox = QHBoxLayout()
        hbox.addWidget(self.btnReture)  # 返回
        hbox.addWidget(self.stockLabel)  # 股票信息
        hbox.addWidget(self.dayLabel)  # 天数
        hbox.addWidget(self.daySel)  # 天数下拉框
        hbox.addWidget(self.mainIndexLabel)  # 主图指标
        hbox.addWidget(self.mainIndexSel)  # 主图指标下拉框
        hbox.addWidget(self.index1Label)  # 指标1
        hbox.addWidget(self.index1Sel)  # 指标1下拉框
        hbox.addWidget(self.index2Label)  # 指标2
        hbox.addWidget(self.index2Sel)  # 指标2下拉框
        hbox.addWidget(self.analyseLabel)  # 分析
        hbox.addWidget(self.analyseSel)  # 分析下拉框
        hbox.addWidget(self.rangeLabel)  # 波段
        hbox.addWidget(self.rangeSel)  # 波段下拉框
        hbox.addWidget(self.btnQuery)  # 查询
        hbox.addStretch(4)

        self.k_plt = pg.PlotWidget()  # 创建一个PyQtGraph.PlotWidget图表对象
        self.k_plt.showGrid(x=True, y=True)
        self.k_plt.setLabel(axis='left', text='价格')  # Y轴标签设置为价格
        self.k_plt.setLabel(axis='bottom', text='日期')  # X轴标签设置为日期

        vbox = QVBoxLayout()
        vbox.addLayout(hbox)
        vbox.addWidget(self.k_plt)
        self.setLayout(vbox)

        # 增加十字线显示
        self.move_slot = pg.SignalProxy(self.k_plt.scene().sigMouseMoved,
                                        rateLimit=50, slot=self.on_pltMouseMove)

        self.daySel.currentIndexChanged.connect(self.daySelChange)  # 切换日期
        self.rangeSel.currentIndexChanged.connect(self.rangeSelChange)  # 切换波段

    # 附加的初始化操作
    def extInit(self):
        # 嵌入其它窗口被调用，使“查询”按钮不可见
        self.btnQuery.setVisible(False)

    # “返回”按钮的槽函数
    def on_btnReture(self):
        if self.parent:  # 如果self.parent存在,说明是嵌入在其他窗口中运行
            self.parent.returnToList()
        else:  # 如果self.parent不存在,说明是独立运行
            self.close()

    # “查询”按钮的槽函数
    # 此处只是为了测试,完整程序中没有这个按钮
    def on_btnQuery(self):
        self.stockQuery('000001.SZ')  # 浦发银行的股票代码

    # 天数下拉框选择不同的日期,自动调用查询
    def daySelChange(self, i):
        self.stockQuery(self.stock_code)

    # 波段下拉框选择不同参数,自动调用查询
    def rangeSelChange(self, i):
        if i == 0:
            self.query_type = "daily"
        elif i == 1:
            self.query_type = "weekly"
        elif i == 2:
            self.query_type = "monthly"
        self.stockQuery(self.stock_code)

    # 查询股票行情信息
    def stockQuery(self, stock_code):
        self.stock_code = stock_code  # 显示当前正在查询的股票: [000001 - 平安银行]
        self.stockLabel.setText("[%s - %s]" % (stock_code[:6], getName(stock_code[:6])))  # 设置显示的标签
        try:
            # 查询的日期由用户自己选择
            date_sel = self.daySel.currentText()[1:-1]
            start_date = datetime.datetime.today() - datetime.timedelta(days=int(date_sel)+1)
            start_date_str = datetime.datetime.strftime(start_date, "%Y%m%d")
            end_date = datetime.datetime.today() - datetime.timedelta(days=1)
            end_date_str = datetime.datetime.strftime(end_date, "%Y%m%d")
            self.displayKPlt(code=stock_code, start=start_date_str, end=end_date_str)
        except Exception as e:
            QMessageBox.information(self, '提示信息', '查询行情信息失败。')

    # 取行情数据并画K线图
    def displayKPlt(self, data=None, code=None, start=None, end=None, queryType='daily'):
        if not data:  # 无数据
            self.data = ts.pro_api().query(self.query_type, ts_code=code, start_date=start, end_date=end).sort_index(ascending=False)
        self.data.set_index(['trade_date'], inplace=True)  # 将日期作为索引
        # 取所有行情的最低价和最高价,用于确定Y轴的范围
        y_min = self.data['low'].min()  # 取所有行情的最低价
        y_max = self.data['high'].max()  # 取所有行情的最高价

        # 循环取每日的序号、开盘价、收盘价、最低价、最高价,组合成元组,加入到元组列表
        data_list = []
        i = 0
        for dates, row in self.data.iterrows():
            # 循环取每日的开盘价、收盘价、最低价、最高价
            open = row['open']
            high = row['high']
            close = row['close']
            low = row['low']
            da = (i, open, close, low, high)  # 组合成元组
            data_list.append(da)  # 加入到元组列表
            i += 1
        # 设置坐标轴显示字符信息
        self.axis_dict = dict(enumerate(self.data.index))
        axis_4 = [(i, list(self.data.index)[i]) for i in
                  range(0, len(self.data.index), 4)]
        axis_16 = [(i, list(self.data.index)[i]) for i in
                  range(0, len(self.data.index), 16)]
        self.k_plt.getAxis("bottom").setTicks([axis_16, axis_4, self.axis_dict.items()])

        self.k_plt.plotItem.clear()  # 清空self.k_plt中所有的图元

        item = CandlesBarItem(data_list)  # 用查询到的行情数据创建一个K线图绘制对象(蜡烛图数据)

        self.k_plt.addItem(item,)  # 将K线图绘制对象加入到self.k_plt中
        self.k_plt.setYRange(y_min, y_max)  # 设定Y轴的取值范围

        # 创建一条水平直线和一条垂直直线,并添加到self.k_plt中
        self.hLine = pg.InfiniteLine(angle=0, movable=False)  # 水平直线
        self.vLine = pg.InfiniteLine(angle=90, movable=False)  # 垂直直线
        self.k_plt.addItem(self.hLine, ignoreBounds=True)
        self.k_plt.addItem(self.vLine, ignoreBounds=True)

        # 创建一个文本项,并将其添加到self.k_plt中,用于浮动显示当前K线的具体信息
        self.label = pg.TextItem()
        self.k_plt.addItem(self.label)

    # 响应鼠标移动事件绘制十字光标
    def on_pltMouseMove(self, event=None):
        if event is None:  # 事件参数为空,则什么都不做
            return
        pos = event[0]  # 从事件获取鼠标当前坐标,这个坐标是场景坐标(获取事件的鼠标位置)
        try:
            if self.k_plt.sceneBoundingRect().contains(pos):  # 如果鼠标坐标在K线图范围内,执行后续操作
                mouse_point = self.k_plt.plotItem.vb.mapSceneToView(pos)  # 场景坐标转换为视图坐标
                index = int(mouse_point.x())  # 通过鼠标所处的X轴坐标,计算K线序号
                if -1 < index < len(self.data.index):  # 如果是合法的K线序号,则执行后续操作
                    self.label.setHtml(  # Html格式的标签文本
                        "<p style='color:white'><strong>日期：{0}</strong></p>"
                        "<p style='color:white'>开盘：{1}</p>"
                        "<p style='color:white'>收盘：{2}</p>"
                        "<p style='color:white'>最高价：<span style='color:red;'>"
                        "{3}</span></p>"
                        "<p style='color:white'>最低价：<span style='color:green;'>"
                        "{4}</span></p>"
                            .format(self.axis_dict[index], self.data['open'][index],
                                    self.data['close'][index], self.data['high'][index],
                                    self.data['low'][index])
                    )
                    self.label.setPos(mouse_point.x(), mouse_point.y())  # 设置图标文本项的位置
                # 设置水平直线和垂直直线的位置,组成十字线
                self.vLine.setPos(mouse_point.x())  # 设置水平直线的位置
                self.hLine.setPos(mouse_point.y())  # 设置垂直直线的位置
        except Exception as e:
            pass


# K线图绘制类
class CandlesBarItem(pg.GraphicsObject):
    def __init__(self, data):
        pg.GraphicsObject.__init__(self)
        # data是一个元组列表,每个元组包含5个数值元素,分别为: 序号、开盘价、收盘价、最低价、最高价
        self.data = data
        self.generatePicture()

    def generatePicture(self):
        self.picture = QtGui.QPicture()  # 创建一个绘图设备self.picture
        p = QtGui.QPainter(self.picture)  # 在self.picture上创建一个QPainter绘图工具p
        p.setPen(pg.mkPen('w'))  # 将绘图工具的画笔设为白色
        w = (self.data[1][0] - self.data[0][0]) / 3.  # 计算K线实体的宽度
        for (i, open, close, low, high) in self.data:  # 取每个元组,每个元组画一个K线
            # print(i, open, close, low, high)
            p.drawLine(QtCore.QPointF(i, low), QtCore.QPointF(i, high))  # 从"最低价"到"最高价"绘制一根K线
            if open > close:  # 根据"开盘价"是否大于"收盘价",来设置画刷的颜色
                p.setBrush(pg.mkBrush('g'))  # 开盘>收盘,绿色
            else:
                p.setBrush(pg.mkBrush('r'))  # 开盘<收盘,红色
            p.drawRect(QtCore.QRectF(i - w, open, w * 2, close - open))  # 画一个矩形作为K线实体
        p.end()  # 结束绘制

    # 重写GraphicsObject类的paint方法
    def paint(self, p, *args):
        # 调用绘图工具的drawPicture方法,在坐标(0,0)处绘制self.picture
        p.drawPicture(0, 0, self.picture)

    # 重写GraphicsObject类的boundingRect方法
    # 该方法在paint方法执行时被调用,paint方法只重新绘制由boundingRect确定的范围
    def boundingRect(self):
        return QtCore.QRectF(self.picture.boundingRect())

    
if __name__ == '__main__':
    app = QApplication(sys.argv)
    w = CandlesBarChart(None)
    w.show()
    sys.exit(app.exec_())